package com.mega.fe.util.klass;

import com.mega.fe.util.MEGA;
import com.mega.fe.asm.ClassReader;
import com.mega.fe.asm.Type;

import java.lang.invoke.MethodHandle;
import java.lang.invoke.MethodHandles;
import java.lang.invoke.MethodType;
import java.lang.invoke.VarHandle;
import java.security.ProtectionDomain;
import java.util.ArrayList;


public class Definer {
    public static boolean isSupportedClassFileVersion(int major, int minor) {
        if (major < 45 || major > 21+44) return false;
        // for major version is between 45 and 55 inclusive, the minor version may be any value
        if (major < 56) return true;
        // otherwise, the minor version must be 0 or 65535
        return minor == 0 || minor == 65535;
    }
    static final int
            NESTMATE_CLASS            = 0x00000001,
            HIDDEN_CLASS              = 0x00000002,
            STRONG_LOADER_LINK        = 0x00000004,
            ACCESS_VM_ANNOTATIONS     = 0x00000008;
    public static class ClassDefiner {
        private final MethodHandles.Lookup lookup;
        private final String name;
        private final byte[] bytes;
        private final int classFlags;

        public ClassDefiner(MethodHandles.Lookup lookup, ClassFile cf, int flags) {
            assert ((flags & HIDDEN_CLASS) != 0 || (flags & STRONG_LOADER_LINK) == STRONG_LOADER_LINK);
            this.lookup = lookup;
            this.bytes = cf.bytes;
            this.name = cf.name;
            this.classFlags = flags;
        }

        public Class<?> defineClass(boolean initialize) throws Throwable {
            return defineClass(initialize, null);
        }
        public Class<?> defineClass(boolean initialize, Object classData) throws Throwable {
            Class<?> lookupClass = lookup.lookupClass();
            ClassLoader loader = lookupClass.getClassLoader() == null ? MEGA.DefineClassLoader.BOOT_LOADER : lookupClass.getClassLoader();
            System.out.printf("!!!ClassDefiner is Getting LookupClassLoader :%s%n", loader);
            MethodHandle mh = MEGA.IMPL_LOOKUP().findStatic(ClassLoader.class, "defineClass0", MethodType.methodType(Class.class
                    , ClassLoader.class, Class.class, String.class, byte[].class, int.class, int.class, ProtectionDomain.class, boolean.class, int.class, Object.class));
            Class<?> c = (Class<?>) mh.invoke( loader, lookupClass, name, bytes, 0, bytes.length, null, initialize, classFlags, classData);
            System.out.println("!!!ClassDefiner is Getting DynamicClass's ClassLoader:" + c.getClassLoader());
            System.out.println("!!!ClassDefiner is Getting DynamicClass's name:" + c.getName());
            System.out.printf("!!!Class:%s defined successfully!", c);
            assert !isNestmate() || c.getNestHost() == lookupClass.getNestHost();
            return c;
        }

        public Class<?> defineClass(Class<?> lookupClass, boolean initialize, Object classData) throws Throwable {
            ClassLoader loader = lookupClass.getClassLoader();
            MethodHandle mh = MEGA.IMPL_LOOKUP().findStatic(ClassLoader.class, "defineClass0", MethodType.methodType(Class.class
                    , ClassLoader.class, Class.class, String.class, byte[].class, int.class, int.class, ProtectionDomain.class, boolean.class, int.class, Object.class));
            Class<?> c = (Class<?>) mh.invoke( loader, lookupClass, name, bytes, 0, bytes.length, null, initialize, classFlags, classData);
            assert !isNestmate() || c.getNestHost() == lookupClass.getNestHost();
            return c;
        }

        private boolean isNestmate() {
            return (classFlags & NESTMATE_CLASS) != 0;
        }
    }
    public static class ClassFile {
        final String name;
        final int accessFlags;
        final byte[] bytes;
        ClassFile(String name, int accessFlags, byte[] bytes) {
            this.name = name;
            this.accessFlags = accessFlags;
            this.bytes = bytes;
        }

        static ClassFile newInstanceNoCheck(String name, byte[] bytes) {
            return new ClassFile(name, 0, bytes);
        }

        /**
         * This method checks the class file version and the structure of `this_class`.
         * and checks if the bytes is a class or interface (ACC_MODULE flag not set)
         * that is in the named package.
         *
         * @throws IllegalArgumentException if ACC_MODULE flag is set in access flags
         * or the class is not in the given package name.
         */
        static ClassFile newInstance(byte[] bytes, String pkgName) {
            int magic = readInt(bytes, 0);
            if (magic != 0xCAFEBABE) {
                throw new ClassFormatError("Incompatible magic value: " + magic);
            }
            int minor = readUnsignedShort(bytes, 4);
            int major = readUnsignedShort(bytes, 6);
            if (!isSupportedClassFileVersion(major, minor)) {
                throw new UnsupportedClassVersionError("Unsupported class file version " + major + "." + minor);
            }

            String name;
            int accessFlags;
            try {
                ClassReader reader = new ClassReader(bytes);
                // ClassReader::getClassName does not check if `this_class` is CONSTANT_Class_info
                // workaround to read `this_class` using readConst and validate the value
                int thisClass = reader.readUnsignedShort(reader.header + 2);
                Object constant = reader.readConst(thisClass, new char[reader.getMaxStringLength()]);
                if (!(constant instanceof Type type)) {
                    throw new ClassFormatError("this_class item: #" + thisClass + " not a CONSTANT_Class_info");
                }
                if (!type.getDescriptor().startsWith("L")) {
                    throw new ClassFormatError("this_class item: #" + thisClass + " not a CONSTANT_Class_info");
                }
                name = type.getClassName();
                accessFlags = reader.readUnsignedShort(reader.header);
            } catch (RuntimeException e) {
                // ASM exceptions are poorly specified
                ClassFormatError cfe = new ClassFormatError();
                cfe.initCause(e);
                throw cfe;
            } 

            // check if it's in the named package
            int index = name.lastIndexOf('.');
            String pn = (index == -1) ? "" : name.substring(0, index); 

            return new ClassFile(name, accessFlags, bytes);
        }

        private static int readInt(byte[] bytes, int offset) {
            if ((offset+4) > bytes.length) {
                throw new ClassFormatError("Invalid ClassFile structure");
            }
            return ((bytes[offset] & 0xFF) << 24)
                    | ((bytes[offset + 1] & 0xFF) << 16)
                    | ((bytes[offset + 2] & 0xFF) << 8)
                    | (bytes[offset + 3] & 0xFF);
        }

        private static int readUnsignedShort(byte[] bytes, int offset) {
            if ((offset+2) > bytes.length) {
                throw new ClassFormatError("Invalid ClassFile structure");
            }
            return ((bytes[offset] & 0xFF) << 8) | (bytes[offset + 1] & 0xFF);
        }
    }
    public static ClassDefiner makeClassDefiner(byte[] bytes) throws Throwable {
        ClassFile cf = ClassFile.newInstance(bytes, MEGA.IMPL_LOOKUP().lookupClass().getPackageName());
        return new ClassDefiner(MEGA.IMPL_LOOKUP(), cf, STRONG_LOADER_LINK);
    }
    public static Class<?> defineClass(byte[] bytes) {
        try {
            return makeClassDefiner(bytes.clone()).defineClass(false);
        } catch (Throwable e) {
            throw new RuntimeException(e);
        }
    }
}
